import numpy as np
from numba import njit
from typing import Callable
# This Python code calculates the drift rate based on the given parameters. Here's a breakdown of what
# each part of the code is doing:
from .drift_function import SSP_drift, DMC_drift, rDMC_drift
import warnings


# congruency int: congruency condition with flanker and target, 1 means congruent, -1 means
@njit
def SSP_trial(theta, congruency=1, dc=1.0, dt=0.001, max_steps=1e5, return_NAN=True):
    """Simulates a trial from the diffusion model."""

    a, ndt, p, sda, rd, z = theta
    n_steps = 0.0
    evidence = a * z

    v_view = []
    trajectory = []

    # Simulate a single DM path
    while evidence > 0 and evidence < a and n_steps < max_steps:
        if congruency == 1:
            drift = p
        else:
            # NOTE: SSP_drift return array, so we need [0]
            drift = SSP_drift(n_steps * dt, p, sda, rd)[0].item()

        trajectory.append(evidence)
        v_view.append(drift)

        # DDM equation
        evidence += drift * dt + np.sqrt(dt) * dc * np.random.normal()

        # Increment step
        n_steps += 1.0

    rt = n_steps * dt + ndt

    if evidence >= a:
        choice = 1
    elif evidence <= 0:
        choice = -1
    elif return_NAN:
        choice = np.nan
        rt = np.nan
    else:
        choice = 1 if evidence>=trajectory[0] else -1


    return rt, choice, v_view, trajectory


@njit
def DMC_trial(theta, congruency=1, dt=0.001, dc=1.0, max_steps=2e5, return_NAN=True):
    """
    DMC_trial simulates a single trial of the drift-diffusion model (DDM) for decision making. 
    Parameters:
    - theta: array-like, model parameters including a, ndt, vc, peak, shape, tau, and z
    - congruency: int, optional, default is 1, congruency condition
    - dt: int, optional, default is 1, time step size
    - dc: float, optional, default is 4.0, diffusion coefficient
    - max_steps: float, optional, default is 2e5, maximum number of iterations

    Returns:
    - rt: float, response time in seconds
    - choice: int, choice made (-1, 0, 1) 
    - v_view: list, drift values over time
    - trajectory: list, evidence accumulation over time
    """

    a, ndt, vc, peak, shape, tau, z = theta

    n_steps = 0.0
    a = a
    tau = tau
    peak = peak
    vc = vc
    evidence = a * z
    congruency = 1 if congruency == 1 else -1
    trajectory = [evidence]
    v_view = [vc]
    
    # Simulate a single DM path
    while evidence > 0 and evidence < a and n_steps < max_steps:

        n_steps += 1.0
        
        # DDM equation
        v = DMC_drift(n_steps*dt*1000, vc, peak * congruency, shape, tau)
        evidence += v * dt + np.sqrt(dt) * dc * np.random.normal()

        v_view.append(v)
        trajectory.append(evidence)

        # Increment step

    rt = n_steps * dt + ndt

    if evidence >= a:
        choice = 1
    elif evidence <= 0:
        choice = -1
    elif return_NAN:
        choice = np.nan
        rt = np.nan
    else:
        choice = 1 if evidence>=trajectory[0] else -1

    return rt, choice, v_view, trajectory



@njit
def rDMC_trial(theta, congruency=1, dt=0.001, dc=0.1, max_steps=1500, return_NAN=True):
    """
    
    Arguments:
    ---
    - theta: list, [a, ndt, A_0, k_c, k_i, d_a, d_c, z, st]
    - congruency: int, 1 for congruent, -1 for incongruent
    - dt: float, time step in seconds
    - dc: float, noise scale
    - max_steps: int, maximum number of steps to simulate
    - return_NAN: bool, if True, return NAN if no response is determined by relative evidence

    Returns:
    ---
    - rt: float, response time in seconds
    - choice: int, choice made (-1, 0, 1) 
    - v_view: list, drift values over time
    - trajectory: list, evidence accumulation over time
    """

    a, ndt, A_0, k_c, k_i, d_a, d_c, z, st = theta

    n_steps = 0.0
    evidence = a * z
    ndt += np.random.normal(0, st)
    trajectory = []
    v_view = []
    
    # Simulate a single DM path
    while evidence > 0 and evidence < a and n_steps < max_steps:

        n_steps += 1.0
        
        # time-depedant drift
        if congruency == 1:
            v = rDMC_drift(n_steps*dt, k_c, A_0, d_a, d_c)
        else: 
            v = rDMC_drift(n_steps*dt, k_i, A_0, -d_a, d_c)
        # ensure v have same scale with dc
        v *= dc * 10
        
        # DDM equation
        evidence += v * dt + np.sqrt(dt) * dc * np.random.normal()

        v_view.append(v)
        trajectory.append(evidence)

    rt = n_steps * dt + ndt

    if evidence >= a:
        choice = 1
    elif evidence <= 0:
        choice = -1
    elif return_NAN:
        choice = np.nan
        rt = np.nan
    else:
        choice = 1 if evidence>=trajectory[0] else -1

    return rt, choice, v_view, trajectory


@njit
def DSTP_trial(theta, congruency=1, dc=1.0, dt=.001, max_steps = 2e3, return_NAN=True):
    """Simulates a trial from the diffusion model."""

    a, ndt, vta, vfl, vss, vp2, ass, zss, z = theta

    n_steps = 0.
    sqrt_st = np.sqrt(dt) * dc
    evidence = a * z
    # initiate the X of start point for stimulus selection
    X_ss = zss * ass
    
    # Drift rate in the first phase
    congruency = 1 if congruency == 1 else -1
    drift = vta + congruency * vfl

    v_view = [drift]
    trajectory = [evidence]

    # Simulate a single DM path
    while (evidence > 0 and evidence < a and n_steps < max_steps):

        noise = np.random.normal(loc=0,scale=1,size = 2)
        # Stimulus selection
        X_ss += vss * dt + noise[0] * sqrt_st
        # Drift rate in the second phase
        if X_ss >= ass:                      # select target
            drift = vp2
            v_view.append(drift)
        elif X_ss <= 0:  # select flanker
            # the flanker have same direction as to target or error
            drift = vp2 if congruency == 1 else -vp2
            v_view.append(drift)

        # DDM equation
        evidence += drift*dt + noise[1] * sqrt_st
        
        trajectory.append(evidence)

        # Increment step
        n_steps += 1.0

    rt = n_steps * dt + ndt
    
    if evidence >= a:
        choice = 1
    elif evidence <= 0:
        choice = -1
    elif return_NAN:
        choice = np.nan
        rt = np.nan
    else:
        choice = 1 if evidence>=trajectory[0] else -1

    return rt, choice, v_view, trajectory



# @njit
# def SSM_trial(theta, dc=1.0, dt=0.001, max_steps=2e5, **kwargs):
#     """Simulates a trial from the diffusion model."""

#     a, ndt, p, sda, rd, z = theta
#     congruency = kwargs.pop("congruency", "None")

#     n_steps = 0.0
#     evidence = a * z

#     v_view = []
#     a_view = []
#     # noise_view = []
#     trajectory = []

#     # Simulate a single DM path
#     while evidence > 0 and evidence < a and n_steps < max_steps:
#         if congruency == 0:
#             drift = SSP_drift(n_steps * dt, p, sda, rd)
#         elif congruency == 1:
#             drift = p

#         # DDM equation
#         evidence += drift * dt + np.sqrt(dt) * dc * np.random.normal()

#         trajectory.append(evidence)
#         v_view.append(drift)

#         # Increment step
#         n_steps += 1.0

#     rt = n_steps * dt + ndt

#     if evidence >= a:
#         choice = 1
#     elif evidence <= 0:
#         choice = -1
#     else:
#         choice = np.nan
#         rt = np.nan


#     return rt, choice, v_view, trajectory
@njit
def DDM_trial(theta, noise_scale=1.0, dt=0.001, max_steps=2e5, return_NAN=True):
    """Simulates a trial from the diffusion model."""

    a, v, t, z, sv, st, sz = theta
    v += np.random.normal(0, sv)
    t += np.random.normal(0, st)
    z += np.random.normal(0, sz)
    n_steps = 0.0
    evidence = a * z

    v_view = [v]
    trajectory = [evidence]
    # Simulate a single DM path
    while evidence > 0 and evidence < a and n_steps < max_steps:

        # DDM equation
        evidence += v * dt + np.sqrt(dt) * noise_scale * np.random.normal()
        trajectory.append(evidence)

        # Increment step
        n_steps += 1.0

    rt = n_steps * dt + t

    if evidence >= a:
        choice = 1
    elif evidence <= 0:
        choice = -1
    elif return_NAN:
        choice = np.nan
        rt = np.nan
    else:
        choice = 1 if evidence>=trajectory[0] else -1

    return rt, choice, v_view, trajectory

@njit
def SEDDM_trial(theta, noise_scale=1.0, dt=0.001, max_steps=2e5, return_NAN=True):
    """Simulates a trial from the diffusion model."""

    a, v, t, z, b, st = theta
    t += np.random.normal(0, st)
    n_steps = 0.0
    evidence = a * z

    v_view = []
    trajectory = []
    # Simulate a single DM path
    while evidence > 0 and evidence < a and n_steps < max_steps:

        v_tmp = v * (1 + b * dt * n_steps)

        # DDM equation
        evidence += v_tmp * dt + np.sqrt(dt) * noise_scale * np.random.normal()
        trajectory.append(evidence)
        v_view.append(v_tmp)


        # Increment step
        n_steps += 1.0

    rt = n_steps * dt + t

    if evidence >= a:
        choice = 1
    elif evidence <= 0:
        choice = -1
    elif return_NAN:
        choice = np.nan
        rt = np.nan
    else:
        choice = 1 if evidence>=trajectory[0] else -1

    return rt, choice, v_view, trajectory

def pointwise_simulator(
    trial_simulator: Callable, prior_theta=None, save_name=None, parallel=False
):

    if save_name is not None:
        import pickle
        from pathlib import Path
        save_name = Path(save_name).with_suffix(".pkl")

        if save_name.exists():
            train_data = pickle.load(open(save_name, 'rb'))
            return train_data

    assert prior_theta is not None, "prior_theta must be provided"
    theta_shape = prior_theta.shape
    n_trials = int(theta_shape[0])

    if parallel:
        from joblib import Parallel, delayed

        def worker(prior_theta):
            while True:
                rt, choice = trial_simulator(prior_theta)[:2]
                if np.isfinite(rt):
                    return (rt, choice)

        out = Parallel(n_jobs=-1)(delayed(worker)(theta) for theta in prior_theta)
        out = np.array(out)
    else:
        counter = 0
        out = np.zeros((n_trials, 2))
        while counter < n_trials:
            rt, choice = trial_simulator(
                prior_theta[:, :-1][counter], prior_theta[:, -1][counter]
            )[:2]
            if np.isfinite(rt):
                out[counter, :] = (rt, choice)
                counter += 1

    train_data = np.column_stack((prior_theta, out))
    if save_name is not None:
        with open(save_name, 'wb') as f:
            pickle.dump(train_data, f)
    return train_data


def batch_simulator(n_samples: int, prior_theta, trial_simulator: Callable):
    counter = 0
    out = np.zeros((n_samples, 2))
    while counter < n_samples:
        rt, choice = trial_simulator(prior_theta)[:2]
        if np.isfinite(rt):
            out[counter, :] = (rt, choice)
            counter += 1

    return out


def experimenti_simulator(
    prior_theta,
    trial_simulator: Callable,
    n_samples: int = 100,
    n_subjects: int = 1,
    sim_kwargs = {}
):
    from functools import partial
    import pandas as pd

    warnings.warn(
        "Call to deprecated function experimenti_simulator. Use dmc_experimenti_simulator instead.",
        category=DeprecationWarning,
        stacklevel=2
    )
    # 初始化一个列表来存储所有被试的数据
    all_data = []
    # 循环遍历每个被试
    for subj_idx in range(n_subjects):
        # 生成一致性条件下的数据
        data_cong = batch_simulator(
            n_samples, prior_theta, partial(trial_simulator, congruency=1, **sim_kwargs)
        )
        data_cong = np.column_stack((data_cong, np.ones(n_samples)))
        # 生成不一致性条件下的数据
        data_incong = batch_simulator(
            n_samples, prior_theta, partial(trial_simulator, congruency=-1, **sim_kwargs)
        )
        data_incong = np.column_stack((data_incong, -np.ones(n_samples)))
        # 将两种条件下的数据合并
        true_data = np.concatenate((data_cong, data_incong))
        # 将被试索引添加到数据中
        true_data = np.column_stack(
            (true_data, np.full((true_data.shape[0], 1), subj_idx))
        )
        # 创建一个 DataFrame 并添加到列表中
        data_df = pd.DataFrame(
            {
                "rt": true_data[:, 0].astype(np.float32),
                "response": true_data[:, 1].astype(np.float32),
                "congruency": true_data[:, 2].astype(np.float32),
                "subj_idx": true_data[:, 3].astype(np.int32)
            }
        )
        all_data.append(data_df)
    # 将所有被试的数据合并到一个大的 DataFrame 中
    combined_data_df = pd.concat(all_data, ignore_index=True)

    return combined_data_df


def dmc_experimenti_simulator(
    prior_theta,
    trial_simulator: Callable,
    n_samples: int = 100,
    n_subjects: int = 1,
    sim_kwargs = {}
):
    return experimenti_simulator(prior_theta, trial_simulator, n_samples, n_subjects, sim_kwargs=sim_kwargs)


# unit test
# n_training_examples = 10
# thetas = SSP_prior(n_training_examples)
# data = model_simulator(prior_theta=thetas, trial_simulator=SSP_trial)
